// 声明包名为 gguf
package llama

// 导入所需的包
import (
	"Lunar-Astral-API/local_server/metadata" // 导入元数据包，用于解析模型元数据
	"fmt"                                    // 格式化输出错误信息
	"log"                                    // 日志记录
	"math"                                   // 提供数学函数（如 Max、Min、Floor 等）
	"os"                                     // 操作系统相关函数（如文件操作、环境变量等）
	"os/exec"                                // 执行系统命令（如 nvidia-smi）
	"strconv"                                // 字符串与数字转换
	"strings"                                // 字符串处理（分割、去除空格等）
)

// getFreeMemory 使用 nvidia-smi 命令查询 GPU 空闲显存，并将结果从 MB 转换为字节返回。
func getFreeMemory() (uint64, error) {
	// 创建执行 nvidia-smi 命令的对象，查询空闲显存并以无表头、无单位的 CSV 格式输出
	cmd := exec.Command("nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits")
	// 执行命令并获取输出结果
	output, err := cmd.CombinedOutput()
	// 若命令执行失败（如未安装 nvidia-smi、驱动异常等），返回错误
	if err != nil {
		return 0, fmt.Errorf("nvidia-smi 执行失败: %w", err)
	}
	// 去除输出结果的前后空白字符并按行分割
	lines := strings.Split(strings.TrimSpace(string(output)), "\n")
	// 如果没有返回任何行数据，说明没有检测到 GPU 或命令无输出
	if len(lines) == 0 {
		return 0, fmt.Errorf("未检测到 GPU 信息")
	}
	// 解析第一行输出为无符号整数，表示空闲显存大小（MB）
	freeMemMB, err := strconv.ParseUint(strings.TrimSpace(lines[0]), 10, 64)
	// 解析失败（如非数字、格式错误）时返回错误
	if err != nil {
		return 0, fmt.Errorf("显存值解析失败: %w", err)
	}
	// 将空闲显存大小从 MB 转换为字节并返回
	return freeMemMB * 1024 * 1024, nil
}

/**
 * @description: 获取模型文件信息和可用显存
 * @param {string} modelPath 模型文件的完整路径
 * @param {string} modelName 模型文件的名称
 * @return {os.FileInfo} 模型文件信息
 * @return {uint64} 可用显存大小（字节）
 * @return {error} 错误信息，若操作过程中无错误则返回 nil
 */
func getFileInfoAndMemory(modelPath, modelName string) (os.FileInfo, uint64, error) {
	// 打印模型加载分隔符
	log.Printf("-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-")
	// 获取模型文件的信息
	fileInfo, err := os.Stat(modelPath)
	// 若获取文件信息失败，记录错误日志并返回错误
	if err != nil {
		log.Printf("GGUF服务[ERROR] -> 模型[%s]加载失败: %v", modelName, err)
		return nil, 0, err
	}
	// 获取可用显存大小
	freeMem, err := getFreeMemory()
	// 若获取显存信息失败，记录错误日志并使用默认值 8GB
	if err != nil {
		log.Printf("GGUF服务[ERROR] -> 显存检测失败，使用默认值 8GB: %v", err)
		freeMem = 8 * 1024 * 1024 * 1024
	} else {
		// 若获取显存信息成功，记录检测到的可用显存大小
		log.Printf("GGUF服务 -> 检测到可用显存: %d MB", freeMem/1024/1024)
	}
	// 返回模型文件信息、可用显存大小和 nil 错误
	return fileInfo, freeMem, nil
}

// 假设 FindFirstMetadataByKeySubstring 返回类型为 int，根据返回值计算最大安全层数
func calculateMetadataLayers(metaData map[string]any, fileSize float64, freeMem uint64) (blockCount, maxSafeLayers int, err error) {
	// 从元数据中获取 .block_count 键对应的值
	val, ok := metadata.FindFirstMetadataByKeySubstring(metaData, ".block_count").(uint32)
	// 检查元数据是否存在
	if !ok {
		return 0, 0, fmt.Errorf("GGUF服务[ERROR] -> 元数据[ block_count : GPU卸载层数 ]不存在")
	}
	// 转换为 int 类型
	blockCount = int(val)
	// 检查 blockCount 是否为 0
	if blockCount == 0 {
		return 0, 0, fmt.Errorf("GGUF服务[ERROR] -> 元数据[ block_count : GPU卸载层数 ]值为 0")
	}
	// 计算每层的大小
	layerSize := fileSize / math.Max(float64(val), 5)
	// 计算最大安全层数，取 1、块数量和 (可用显存 * 0.65 / 每层大小) 中的中间值
	maxSafeLayers = int(math.Max(1, math.Min(float64(val), float64(freeMem)*0.65/layerSize)))
	// 返回计算得到的最大安全层数
	return blockCount, maxSafeLayers, nil
}

func getMetadataLayersAndMemory(modelPath, modelName string, metadata map[string]interface{}) (blockCount, maxSafeLayers int, err error) {
	fileInfo, freeMem, err := getFileInfoAndMemory(modelPath, modelName)
	if err != nil {
		return 0, 0, err
	}
	// 根据模型元数据、文件大小和可用显存计算块数量和最大安全 GPU 加速层级
	return calculateMetadataLayers(metadata, float64(fileInfo.Size()), freeMem)
}
